#!/usr/bin/bash

IMAGE_INFO="/usr/share/ublue-os/image-info.json"
IMAGE_NAME=$(jq -r '."image-name"' < $IMAGE_INFO)
IMAGE_BRANCH=$(jq -r '."image-branch"' < $IMAGE_INFO)
FEDORA_VERSION=$(jq -r '."fedora-version"' < $IMAGE_INFO)

# SCRIPT VERSION
HWS_VER=64
HWS_VER_FILE="/etc/bazzite/hws_version"
HWS_VER_RAN=$(cat $HWS_VER_FILE)

# IMAGE IDENTIFIERS
KNOWN_IMAGE_NAME_FILE="/etc/bazzite/image_name"
KNOWN_IMAGE_NAME=$(cat $KNOWN_IMAGE_NAME_FILE)
KNOWN_IMAGE_BRANCH_FILE="/etc/bazzite/image_branch"
KNOWN_IMAGE_BRANCH=$(cat $KNOWN_IMAGE_BRANCH_FILE)
KNOWN_FEDORA_VERSION_FILE="/etc/bazzite/fedora_version"
KNOWN_FEDORA_VERSION=$(cat $KNOWN_FEDORA_VERSION_FILE)

if [[ "$KNOWN_IMAGE_NAME" == "$IMAGE_NAME" && \
    "$KNOWN_IMAGE_BRANCH" == "$IMAGE_BRANCH" && \
    "$KNOWN_FEDORA_VERSION" == "$FEDORA_VERSION" && \
    "$HWS_VER_RAN" == "$HWS_VER" ]]; then
  echo "Hardware setup already ran for version $HWS_VER_RAN, skipping"
  exit 0
fi

# GLOBAL
SYS_ID="$(/usr/libexec/hwsupport/sysid)"
VEN_ID="$(cat /sys/devices/virtual/dmi/id/chassis_vendor)"
CPU_VENDOR=$(grep "vendor_id" "/proc/cpuinfo" | uniq | awk -F": " '{ print $2 }')
CPU_MODEL=$(grep "model name" "/proc/cpuinfo" | uniq | awk -F": " '{ print $2 }')
MINIMUM_FREE_ZRAM=$(awk '/MemTotal/ {printf "%.0f", $2 * 0.01}' /proc/meminfo)
CURRENT_FREE_ZRAM=$(sysctl vm.min_free_kbytes | awk '{print $3}')
KARGS=$(rpm-ostree kargs)
NEEDED_KARGS=()

# Add directory for fixups
mkdir -p /etc/bazzite/fixups

# KERNEL ARGUMENTS
echo "Current kargs: $KARGS"

if /usr/libexec/hwsupport/valve-hardware; then
  echo "Checking for needed karg changes (Jupiter/Galileo)"

  if [[ ! $KARGS =~ "amd_iommu" ]]; then
    NEEDED_KARGS+=("--append-if-missing=amd_iommu=off")
  fi
fi

if [[ $IMAGE_NAME =~ "deck" && ":AOKZOE A1 AR07:" =~ ":$SYS_ID:" ]]; then
  echo "AOKZOE A1 on deck build detected, fixing edid"
  if [[ ! $KARGS =~ "drm.edid_firmware" ]]; then
    NEEDED_KARGS+=("--append-if-missing=drm.edid_firmware=eDP-1:edid/aokzoea1ar07_edid.bin")
  fi
fi

if [[ ":Framework:" =~ ":$VEN_ID:" ]]; then
  if [[ "$SYS_ID" == "Laptop 16 ("* ]]; then
    if [[ $KARGS =~ "amdgpu.dcdebugmask" ]]; then
      echo "Removing old, no longer needed Framework 16 fix."
      NEEDED_KARGS+=("--delete-if-present=amdgpu.dcdebugmask=0x410")
    fi
  fi
  if [[ "GenuineIntel" == "$CPU_VENDOR" ]]; then
    if [[ ! $KARGS =~ "hid_sensor_hub" ]]; then
      echo "Intel Framework Laptop detected, applying needed keyboard fix"
      NEEDED_KARGS+=("--append-if-missing=module_blacklist=hid_sensor_hub");
    fi
  fi
fi

#
# GTT SIZE cleanup
#
if [ ! -f /etc/bazzite/fixups/gttsize ]; then
  echo "Removing gttsize for device in kernel if exists (default is 50%)"
  if [[ ":83E1:" =~ ":$SYS_ID:" ]]; then
    NEEDED_KARGS+=("--delete-if-present=amdgpu.gttsize=8128")
  elif [[ ":ROG Ally RC71L:" =~ ":$SYS_ID:" ]]; then
    NEEDED_KARGS+=("--delete-if-present=amdgpu.gttsize=8128")
  elif [[ ":ROG Ally X RC72LA:" =~ ":$SYS_ID:" ]]; then
    NEEDED_KARGS+=("--delete-if-present=amdgpu.gttsize=12192")
  elif [[ ":ONEXPLAYER F1:ONEXPLAYER F1L:ONEXPLAYER F1 EVA-01:ONEXPLAYER X1Pro:Loki Max:Zeenix Pro:" =~ ":$SYS_ID:" ]]; then
    NEEDED_KARGS+=("--delete-if-present=amdgpu.gttsize=16256")
  elif /usr/libexec/hwsupport/valve-hardware; then
    NEEDED_KARGS+=("--delete-if-present=amdgpu.gttsize=16256")
    NEEDED_KARGS+=("--delete-if-present=amdgpu.gttsize=8128")
  fi
fi

# Remove sg_display
if [[ ":83E1:ROG Ally RC71L:ROG Ally X RC72LA:" =~ ":$SYS_ID:" ]]; then
  if [[ ! -f /etc/bazzite/fixups/sg_display2 && $KARGS =~ "amdgpu.sg_display" ]]; then
    echo "Removing sd_display=0 (it is not needed anymore)"
    NEEDED_KARGS+=("--delete-if-present=amdgpu.sg_display=0")
  fi
fi

if [[ ! -f /etc/bazzite/fixups/orientation_old2 ]]; then
  echo "Removing panel orientation for device in kernel if exists"
  if [[ ":ROG Ally RC71L:" =~ ":$SYS_ID:" ]]; then
    NEEDED_KARGS+=("--delete-if-present=video=eDP-1:panel_orientation=left_side_up")
  elif [[ ":G1617-01:" =~ ":$SYS_ID:" && ! $CPU_MODEL =~ 8640U|8840U|7640U ]]; then
    NEEDED_KARGS+=("--delete-if-present=video=eDP-1:panel_orientation=right_side_up")
  elif [[ ":G1617-01:" =~ ":$SYS_ID:" ]]; then
    NEEDED_KARGS+=("--delete-if-present=video=eDP-1:panel_orientation=right_side_up")
  elif [[ ":WIN2:" =~ ":$SYS_ID:" ]]; then
    NEEDED_KARGS+=("--delete-if-present=video=eDP-1:panel_orientation=right_side_up")
  elif [[ ":ONEXPLAYER F1:ONEXPLAYER F1L:ONEXPLAYER F1 EVA-01:ONEXPLAYER X1Pro:Loki Max:Zeenix Pro:" =~ ":$SYS_ID:" ]]; then
    NEEDED_KARGS+=("--delete-if-present=video=eDP-1:panel_orientation=left_side_up")
  fi
fi

# Slide reboots randomly, acpi=strict seems to fix it
if [[ ":SLIDE:" =~ ":$SYS_ID:" ]]; then
  if [[ ! $KARGS =~ "acpi=strict" ]]; then
    echo "Adding acpi=strict for Slide to avoid random reboots"
    NEEDED_KARGS+=("--append-if-missing=acpi=strict")
  fi
fi

if [[ ":ROG Flow Z13 GZ302EA:" =~ "$:SYS_ID:" ]]; then
  if [[ ! $KARGS =~ "amdgpu.dcdebugmask" ]]; then
    echo "ROG Flow Z13 2025, fixing display artifacts"
    NEEDED_KARGS+=("--append-if-missing=amdgpu.dcdebugmask=0x410")
  fi
fi

if [[ ! $KARGS =~ "disable_ertm" ]]; then
  echo "Adding needed kargs for bluetooth"
  NEEDED_KARGS+=("--append-if-missing=bluetooth.disable_ertm=1")
fi

if [[ ! -f /etc/bazzite/fixups/preempt_cleanup2 && $KARGS =~ "preempt" ]]; then
  echo "Removing full preemption (once)"
  NEEDED_KARGS+=("--delete-if-present=preempt=full")
fi

if [[ $KARGS =~ "nomodeset" ]]; then
  echo "Removing nomodeset"
  NEEDED_KARGS+=("--delete-if-present=nomodeset")
fi

if [[ $IMAGE_NAME =~ "deck" || $IMAGE_NAME =~ "ally" ]]; then
  if /usr/libexec/hwsupport/valve-hardware; then
    if [[ ":Jupiter:" =~ ":$SYS_ID:" ]]; then
      RESOLUTION=$(lshw -json -c display | jq -r .[]."configuration"."resolution")
      if [[ "${RESOLUTION}" = "1200,1920" ]]; then
        echo "Jupiter with DeckHD detected, disabling BIOS updates..."
        systemctl disable --now jupiter-biosupdate.service
      elif [[ "$(awk '/MemTotal/{print $(NF-1)}' /proc/meminfo)" == "31664740" ]]; then
        echo "Jupiter with 32GB RAM detected, disabling BIOS updates..."
        systemctl disable --now jupiter-biosupdate.service
      fi
    fi

    systemctl enable --now jupiter-fan-control.service
    systemctl enable --now vpower.service
  else
    echo "Generic device detected. Performing setup..."
    systemctl disable --now jupiter-fan-control.service
    systemctl disable --now vpower.service
    systemctl disable --now jupiter-biosupdate.service
    systemctl disable --now jupiter-controller-update.service
    systemctl disable --now ryzenadj.service
  fi
  systemctl enable --now ds-inhibit.service
fi

# ZRAM MINIMUM-FREE CONFIGURATION
echo "Current minimum-free ZRAM value: $CURRENT_FREE_ZRAM"

if ((MINIMUM_FREE_ZRAM > CURRENT_FREE_ZRAM)); then
    sysctl -w "vm.min_free_kbytes=${MINIMUM_FREE_ZRAM}"
    echo "Found needed minimum-free ZRAM changes, applying the following: ${MINIMUM_FREE_ZRAM}"
else
  echo "No minimum-free ZRAM changes needed"
fi

# FRAMEWORK AMD FIXES
if [[ ":Framework:" =~ ":$VEN_ID:" ]]; then
  if [[ "AuthenticAMD" == "$CPU_VENDOR" ]]; then
    # Remove ryzen_smu module load if it exists
    rm -rf /etc/modules-load.d/ryzen_smu.conf

    # Framework 13 AMD fixes
    if [[ $SYS_ID == "Laptop ("* ]]; then
      if [[ ! -f /etc/modprobe.d/alsa.conf ]]; then
        echo 'Fixing 3.5mm jack'
        echo "options snd-hda-intel index=1,0 model=auto,dell-headset-multi" > /etc/modprobe.d/alsa.conf
        echo 0 > /sys/module/snd_hda_intel/parameters/power_save
      fi
      if [[ ! -f /etc/udev/rules.d/20-suspend-fixes.rules ]]; then
        echo 'Fixing suspend issue'
        echo "ACTION==\"add\", SUBSYSTEM==\"serio\", DRIVERS==\"atkbd\", ATTR{power/wakeup}=\"disabled\"" > /etc/udev/rules.d/20-suspend-fixes.rules
      fi
    fi
  fi
fi

# SURFACE FIXES
if [[ ":Microsoft:" =~ ":$VEN_ID:" ]]; then
  if ! [[ -f "/etc/modules-load.d/surface.conf" ]]; then
    if [[ "AuthenticAMD" == "$CPU_VENDOR" ]]; then
      printf "# Support GPIO pins on AMD Surfaces\npinctrl_amd\n" > "/etc/modules-load.d/surface.conf"
    elif [[ "GenuineIntel" == "$CPU_VENDOR" ]]; then
      sed 's/^ *//' > "/etc/modules-load.d/surface.conf" << EOF
        # Support GPIO pins on Intel Surfaces
        # Surface Book 2
        pinctrl_sunrisepoint

        # For Surface Laptop 3/Surface Book 3
        pinctrl_icelake

        # For Surface Laptop 4/Surface Laptop Studio
        pinctrl_tigerlake

        # For Surface Pro 9/Surface Laptop 5
        pinctrl_alderlake

        # For Surface Pro 10/Surface Laptop 6
        pinctrl_meteorlake

        # Support LPSS on Intel Surfaces
        intel_lpss
        intel_lpss_pci

EOF
    fi
    sed 's/^ *//' >> "/etc/modules-load.d/surface.conf" << EOF
      # Add modules necessary for Disk Encryption via keyboard
      surface_aggregator
      surface_aggregator_registry
      surface_aggregator_hub
      surface_hid_core
      8250_dw

      # Surface Laptop 3/Surface Book 3 and later
      surface_hid
      surface_kbd
EOF
  fi
fi

# WAYDROID FIX
if [[ -f "/var/lib/waydroid/lxc/waydroid/config" ]]; then
  echo "Removing unneeded apparmor entry from Waydroid LXC"
  sed -i '/lxc\.apparmor\.profile\s*=\s*unconfined/d' "/var/lib/waydroid/lxc/waydroid/config"
fi

# HOSTNAME FIX
# If the hostname is too long Distrobox will fail during setup
# Let's check the length and reset it to something sensible if that happens.
if (( $(hostname | wc -m) > 20 )); then
  hostnamectl set-hostname bazzite
fi

# PPD FIX
# Correct broken config files on some systems
if [[ -f "/etc/tuned/ppd.conf" ]] && ! grep "\[battery\]" "/etc/tuned/ppd.conf"; then
  cp -f "/usr/etc/tuned/ppd.conf" "/etc/tuned/ppd.conf"
fi

# Set default target to graphical, fixes rebase from base image
if grep -qv "graphical.target" <<< "$(systemctl get-default)"; then
  systemctl set-default graphical.target
fi

# Cleanup btrfs oopsie
# default config has a 10 year limit, which is ridiculous
# FIXME: This fix was added in 12/2024, remove it in a few months
if [ ! -f /etc/bazzite/fixups/snapper_cleanup ] &&
  [ -f /etc/snapper/configs/root ] &&
  [ $(snapper --no-headers --csvout -c root get-config --columns key,value | grep "TIMELINE_LIMIT_YEARLY,10" ) ] &&
  [ $(snapper --no-headers --csvout -c root get-config --columns key,value | grep "SUBVOLUME,/var/home" ) ]; then
  echo "DELETING SNAPPER SNAPSHOTS CREATED BY A PREVIOUS UPDATE"
  echo "THIS IS FOR /var/home. IF YOU INTENDED THIS CHANGE, "
  echo "RECREATE YOUR CONFIG AND MAKE SURE TIMELINE_LIMIT_YEARLY=\"10\""
  echo "IS SET TO SOMETHING DIFFERENT."
  SNAPSHOTS=$(snapper --no-headers --csvout --config root list --columns number,description | grep timeline | awk -F, '{print $1}')
  if [ -n "$SNAPSHOTS" ]; then
    snapper -c root delete $SNAPSHOTS
  fi
  snapper -c root delete-config
fi
touch /etc/bazzite/fixups/snapper_cleanup

if [ ! -f /etc/bazzite/fixups/grub_fixup ]; then
  # Grub settings
  # Set grub timeout to 3 seconds
  if [ -z $(grep "GRUB_TIMEOUT=" /etc/default/grub) ]; then
    echo "GRUB_TIMEOUT=3" >> /etc/default/grub
  else
    sed -i 's/GRUB_TIMEOUT=[0-9]*/GRUB_TIMEOUT=3/' /etc/default/grub
  fi

  # Nuke hidden timeout as it will cause issues if GRUB_TIMEOUT exists
  if [ -n $(grep "GRUB_HIDDEN_TIMEOUT=" /etc/default/grub) ]; then
    sed -i 's/GRUB_HIDDEN_TIMEOUT=[0-9]*//' /etc/default/grub
  fi

  # set timeout style to menu to show grub during countdown
  # (unless GRUB_TIMEOUT=0)
  if [ -z $(grep "GRUB_TIMEOUT_STYLE=" /etc/default/grub) ]; then
    echo "GRUB_TIMEOUT_STYLE=menu" >> /etc/default/grub
  else
    sed -i 's/GRUB_TIMEOUT_STYLE=.*/GRUB_TIMEOUT_STYLE=menu/' /etc/default/grub
  fi

  # Autohide grub menu if boot is successful
  # menu_auto_hide=2 hides grub menu if previous boot is successful
  # menu_auto_hide=1 same, except not if Windows/Other OS is installed
  # Deck images cant use grub menu and often dual boot, so hide it always
  # For desktop images, hide it only if no other OS is installed
  if [[ $IMAGE_NAME =~ "deck" ]]; then
    grub2-editenv - set menu_auto_hide=2
  else
    grub2-editenv - set menu_auto_hide=1
  fi

  # Update grub
  if [ -d /sys/firmware/efi ]; then
    echo "Updating EFI grub config"
    grub2-mkconfig -o /etc/grub2-efi.cfg
  else
    echo "Updating BIOS grub config"
    grub2-mkconfig -o /etc/grub2.cfg
  fi
fi
touch /etc/bazzite/fixups/grub_fixup

if [ ! -f /etc/bazzite/fixups/grub_blank ] &&
    [ -z $(grep "GRUB_TERMINAL_OUTPUT=" /etc/default/grub) ]; then
  echo 'GRUB_TERMINAL_OUTPUT="console"' >> /etc/default/grub

  # Update grub
  if [ -d /sys/firmware/efi ]; then
    echo "Updating EFI grub config"
    grub2-mkconfig -o /etc/grub2-efi.cfg
  else
    echo "Updating BIOS grub config"
    grub2-mkconfig -o /etc/grub2.cfg
  fi
fi
touch /etc/bazzite/fixups/grub_blank

# Fix missing hostname file breaking distrobox
if [ ! -f /etc/hostname ]; then
  touch /etc/hostname
fi

# Remove old hhd@ enables
# Check if we've already removed old hhd@ service enablements
if [ ! -f /etc/bazzite/fixups/hhd_service_cleanup ]; then
  echo "Cleaning up old hhd@ service enables"
  for service in $(systemctl list-units --type=service --all 'hhd@*.service' --no-legend | awk '{print $1}'); do
    echo "Disabling $service"
    systemctl disable --now "$service" || true
  done
  touch /etc/bazzite/fixups/hhd_service_cleanup
fi

if [[ -n "$NEEDED_KARGS" ]]; then
  echo "Found needed karg changes, applying the following: ${NEEDED_KARGS[*]}"
  plymouth display-message --text="Preparing System - Please wait, this may take a while" || true
  rpm-ostree kargs ${NEEDED_KARGS[*]} || exit 1
  SHOULD_REBOOT=true
else
  echo "No karg changes needed"
fi
# Applied / did not need to apply
touch /etc/bazzite/fixups/sg_display2
touch /etc/bazzite/fixups/preempt_cleanup2
touch /etc/bazzite/fixups/orientation_old2
touch /etc/bazzite/fixups/gttsize

mkdir -p /etc/bazzite
echo $HWS_VER > $HWS_VER_FILE
echo $IMAGE_NAME > $KNOWN_IMAGE_NAME_FILE
echo $FEDORA_VERSION > $KNOWN_FEDORA_VERSION_FILE
echo $IMAGE_BRANCH > $KNOWN_IMAGE_BRANCH_FILE

if [[ "$SHOULD_REBOOT" == true ]]; then
  echo "Rebooting to apply karg changes"
  reboot
fi
